/*
 * Copyright (c) 2011-2022, baomidou (jobob@qq.com).
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.sky.wsp.mybatis.plus.plugins.inner;

import com.baomidou.mybatisplus.core.plugins.InterceptorIgnoreHelper;
import com.baomidou.mybatisplus.core.toolkit.*;
import com.baomidou.mybatisplus.extension.plugins.inner.BaseMultiTableInnerInterceptor;
import com.baomidou.mybatisplus.extension.plugins.inner.InnerInterceptor;
import com.baomidou.mybatisplus.extension.toolkit.PropertyMapper;
import com.sky.wsp.mybatis.plus.plugins.handler.MyDataPermissionHandler;
import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.expression.LongValue;
import net.sf.jsqlparser.expression.Parenthesis;
import net.sf.jsqlparser.expression.operators.conditional.OrExpression;
import net.sf.jsqlparser.expression.operators.relational.EqualsTo;
import net.sf.jsqlparser.expression.operators.relational.ExpressionList;
import net.sf.jsqlparser.expression.operators.relational.InExpression;
import net.sf.jsqlparser.schema.Column;
import net.sf.jsqlparser.schema.Table;
import net.sf.jsqlparser.statement.delete.Delete;
import net.sf.jsqlparser.statement.insert.Insert;
import net.sf.jsqlparser.statement.select.*;
import net.sf.jsqlparser.statement.update.Update;
import net.sf.jsqlparser.statement.update.UpdateSet;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.ParameterMapping;
import org.apache.ibatis.mapping.SqlCommandType;
import org.apache.ibatis.reflection.MetaObject;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;

import java.sql.Connection;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.List;
import java.util.Properties;
import java.util.stream.Collectors;

/**
 * 参考MybatisPlus（版本3.4.0）的多租户拦截器插件，改造成本数据权限拦截器插件
 *
 * @author wangshaopeng@talkweb.com.cn
 * @Date 2023-01-11
 * @see com.baomidou.mybatisplus.extension.plugins.inner.TenantLineInnerInterceptor
 */
@SuppressWarnings({"rawtypes"})
public class MyDataPermissionInterceptor extends BaseMultiTableInnerInterceptor implements InnerInterceptor {

    private MyDataPermissionHandler dataPermissionHandler;

    @Override
    public void beforeQuery(Executor executor, MappedStatement ms, Object parameter, RowBounds rowBounds, ResultHandler resultHandler, BoundSql boundSql) throws SQLException {
        if (InterceptorIgnoreHelper.willIgnoreDataPermission(ms.getId())) {
            return;
        }
        PluginUtils.MPBoundSql mpBs = PluginUtils.mpBoundSql(boundSql);
        mpBs.sql(parserSingle(mpBs.sql(), null));
    }

    @Override
    public void beforePrepare(StatementHandler sh, Connection connection, Integer transactionTimeout) {
        PluginUtils.MPStatementHandler mpSh = PluginUtils.mpStatementHandler(sh);
        MappedStatement ms = mpSh.mappedStatement();
        SqlCommandType sct = ms.getSqlCommandType();
        if (sct == SqlCommandType.INSERT || sct == SqlCommandType.UPDATE || sct == SqlCommandType.DELETE) {
            if (InterceptorIgnoreHelper.willIgnoreDataPermission(ms.getId())) {
                return;
            }
            // 1、与statement的columns对应的参数properties
            PluginUtils.MPBoundSql mpBs = mpSh.mPBoundSql();
            List<ParameterMapping> parameterMappings = mpBs.parameterMappings();

            // 2、将请求参数VO映射成反射代理对象，以便反射出属性值
            Object parameterObject = sh.getParameterHandler().getParameterObject();
            // TODO 如果入参不是VO，而是List(批量操作）、Map等，该如何处理？
//            if (parameterObject instanceof Map) {
//
//            }else if (parameterObject instanceof List){
//
//            }
            MetaObject metaParameters = mpSh.configuration().newMetaObject(parameterObject);

            DataPermissionParameterVO vo = new DataPermissionParameterVO();
            vo.setMetaParameters(metaParameters);
            vo.setParameterMappings(parameterMappings);

            mpBs.sql(parserMulti(mpBs.sql(), vo));
        }
    }

    @Override
    protected void processSelect(Select select, int index, String sql, Object obj) {
        final String whereSegment = (String) obj;
        processSelectBody(select.getSelectBody(), whereSegment);
        List<WithItem> withItemsList = select.getWithItemsList();
        if (!CollectionUtils.isEmpty(withItemsList)) {
            withItemsList.forEach(withItem -> processSelectBody(withItem, whereSegment));
        }
    }

    @Override
    protected void processInsert(Insert insert, int index, String sql, Object obj) {
        if (dataPermissionHandler.ignoreTable(insert.getTable().getName())) {
            // 过滤退出执行
            return;
        }
        List<Column> columns = insert.getColumns();
        if (CollectionUtils.isEmpty(columns)) {
            // 针对不给列名的insert 不处理
            return;
        }

        DataPermissionParameterVO vo = (DataPermissionParameterVO) obj;
        MetaObject metaParameters = vo.getMetaParameters();

        // 1、找到column下标
        String dataPermissionIdColumn = dataPermissionHandler.getDataPermissionIdColumn();
        int columnIndex = this.getColumnIndex(columns, dataPermissionIdColumn);

        // 2、根据column下标，找到entity的property
        List<ParameterMapping> parameterMappings = vo.getParameterMappings();
        ParameterMapping parameterMapping = parameterMappings.get(columnIndex);
        String property = parameterMapping.getProperty();

        // 3、根据entity的property，从请求参数中反射出数据权限的值
        Long insertOrgId = (Long) metaParameters.getValue(property);

        // 4、判断当前账号，是否有权限操作请求参数中的数据权限id
        boolean contains = dataPermissionHandler.getDataPermissionIdSet().contains(insertOrgId);
        if (!contains) {
            throw ExceptionUtils.mpe("data permission denied before insert！");
        }

//        columns.add(new Column(dataPermissionIdColumn));
//
//        // fixed gitee pulls/141 duplicate update
//        List<Expression> duplicateUpdateColumns = insert.getDuplicateUpdateExpressionList();
//        if (CollectionUtils.isNotEmpty(duplicateUpdateColumns)) {
//            EqualsTo equalsTo = new EqualsTo();
//            equalsTo.setLeftExpression(new StringValue(dataPermissionIdColumn));
////            equalsTo.setRightExpression(dataPermissionHandler.getDataPermissionIdSet());
//            duplicateUpdateColumns.add(equalsTo);
//        }
//
//        Select select = insert.getSelect();
//        if (select != null) {
//            this.processInsertSelect(select.getSelectBody(), (String) obj);
//        } else if (insert.getItemsList() != null) {
//            // fixed github pull/295
//            ItemsList itemsList = insert.getItemsList();
//            Expression dataPermissionIdSet = dataPermissionHandler.getDataPermissionIdSet();
//            if (itemsList instanceof MultiExpressionList) {
//                ((MultiExpressionList) itemsList).getExpressionLists().forEach(el -> el.getExpressions().add(dataPermissionIdSet));
//            } else {
//                ((ExpressionList) itemsList).getExpressions().add(dataPermissionIdSet);
//            }
//        } else {
//            throw ExceptionUtils.mpe("Failed to process multiple-table update, please exclude the tableName or statementId");
//        }
    }

    protected int getColumnIndex(List<Column> columns, String dataPermissionIdColumn) {
        int i = 0;
        for (Column column : columns) {
            if (column.getColumnName().equalsIgnoreCase(dataPermissionIdColumn)) {
                return i;
            }
            ++i;
        }
        return -1;
    }

    /**
     * update 语句处理
     */
    @Override
    protected void processUpdate(Update update, int index, String sql, Object obj) {
        final Table table = update.getTable();
        if (dataPermissionHandler.ignoreTable(table.getName())) {
            // 过滤退出执行
            return;
        }
        String dataPermissionIdColumn = dataPermissionHandler.getDataPermissionIdColumn();
        DataPermissionParameterVO vo = (DataPermissionParameterVO) obj;
        MetaObject metaParameters = vo.getMetaParameters();

        // 1、找到column下标
        ArrayList<UpdateSet> updateSets = update.getUpdateSets();
        int columnIndex = this.getUpdateColumnIndex(updateSets, dataPermissionIdColumn);

        // 2、根据column下标，找到entity的property
        List<ParameterMapping> parameterMappings = vo.getParameterMappings();
        ParameterMapping parameterMapping = parameterMappings.get(columnIndex);
        String property = parameterMapping.getProperty();

        // 3、根据entity的property，从请求参数中反射出数据权限的值
        Long insertOrgId = (Long) metaParameters.getValue(property);

        // 4、判断当前账号，是否有权限操作请求参数中的数据权限id
        boolean contains = dataPermissionHandler.getDataPermissionIdSet().contains(insertOrgId);
        if (!contains) {
            throw ExceptionUtils.mpe("data permission denied before update！");
        }

        update.setWhere(this.andExpression(table, update.getWhere(), null));
    }

    protected int getUpdateColumnIndex(ArrayList<UpdateSet> updateSets, String dataPermissionIdColumn) {
        int i = 0;
        for (UpdateSet updateSet : updateSets) {
            // TODO 这里直接取0下标
            Column column = updateSet.getColumns().get(0);
            if (column.getColumnName().equalsIgnoreCase(dataPermissionIdColumn)) {
                return i;
            }
            ++i;
        }
        return -1;
    }

    /**
     * delete 语句处理
     */
    @Override
    protected void processDelete(Delete delete, int index, String sql, Object obj) {
        if (dataPermissionHandler.ignoreTable(delete.getTable().getName())) {
            // 过滤退出执行
            return;
        }
        delete.setWhere(this.andExpression(delete.getTable(), delete.getWhere(), (String) obj));
    }

    /**
     * 处理 insert into select
     * <p>
     * 进入这里表示需要 insert 的表启用了数据权限,则 select 的表都启动了
     *
     * @param selectBody SelectBody
     */
    protected void processInsertSelect(SelectBody selectBody, final String whereSegment) {
        PlainSelect plainSelect = (PlainSelect) selectBody;
        FromItem fromItem = plainSelect.getFromItem();
        if (fromItem instanceof Table) {
            // fixed gitee pulls/141 duplicate update
            processPlainSelect(plainSelect, whereSegment);
            appendSelectItem(plainSelect.getSelectItems());
        } else if (fromItem instanceof SubSelect) {
            SubSelect subSelect = (SubSelect) fromItem;
            appendSelectItem(plainSelect.getSelectItems());
            processInsertSelect(subSelect.getSelectBody(), whereSegment);
        }
    }

    /**
     * 追加 SelectItem
     *
     * @param selectItems SelectItem
     */
    protected void appendSelectItem(List<SelectItem> selectItems) {
        if (CollectionUtils.isEmpty(selectItems)) {
            return;
        }
        if (selectItems.size() == 1) {
            SelectItem item = selectItems.get(0);
            if (item instanceof AllColumns || item instanceof AllTableColumns) {
                return;
            }
        }
        selectItems.add(new SelectExpressionItem(new Column(dataPermissionHandler.getDataPermissionIdColumn())));
    }

    /**
     * 数据权限字段别名设置
     * <p>orgId 或 tableAlias.orgId</p>
     *
     * @param table 表对象
     * @return 字段
     */
    protected Column getAliasColumn(Table table) {
        StringBuilder column = new StringBuilder();
        // todo 该起别名就要起别名,禁止修改此处逻辑
        if (table.getAlias() != null) {
            column.append(table.getAlias().getName()).append(StringPool.DOT);
        }
        column.append(dataPermissionHandler.getDataPermissionIdColumn());
        return new Column(column.toString());
    }

    @Override
    public void setProperties(Properties properties) {
        PropertyMapper.newInstance(properties).whenNotBlank("myDataPermissionHandler",
                ClassUtils::newInstance, this::setDataPermissionHandler);
    }

    /**
     * 构建数据权限条件表达式
     *
     * @param table        表对象
     * @param where        当前where条件
     * @param whereSegment 所属Mapper对象全路径
     * @return 数据权限条件表达式
     * @see BaseMultiTableInnerInterceptor#buildTableExpression(Table, Expression, String)
     */
    @Override
    public Expression buildTableExpression(final Table table, final Expression where, final String whereSegment) {
        if (dataPermissionHandler.ignoreTable(table.getName())) {
            return null;
        }

        EqualsTo equalsTo = new EqualsTo(new Column(dataPermissionHandler.getDataCreatorColumn()), new LongValue(dataPermissionHandler.getDataCreator()));

        List<Long> set = dataPermissionHandler.getDataPermissionIdSet();
        List<Expression> dataPermissionSet = set.stream()
                .map(id -> new LongValue(id))
                .collect(Collectors.toList());
        ExpressionList expressionList = new ExpressionList(dataPermissionSet);
        InExpression inExpression = new InExpression(getAliasColumn(table), expressionList);

        OrExpression orExpression = new OrExpression();
        orExpression.setLeftExpression(equalsTo);
        orExpression.setRightExpression(inExpression);
        return new Parenthesis(orExpression);
    }

    public MyDataPermissionInterceptor() {
    }

    public MyDataPermissionInterceptor(MyDataPermissionHandler dataPermissionHandler) {
        this.dataPermissionHandler = dataPermissionHandler;
    }

    public MyDataPermissionHandler getDataPermissionHandler() {
        return dataPermissionHandler;
    }

    public void setDataPermissionHandler(MyDataPermissionHandler dataPermissionHandler) {
        this.dataPermissionHandler = dataPermissionHandler;
    }
}
